# RAG高级---句子窗口检索

import os
import openai
from IPython.utils import docs
from llama_index import SimpleDirectoryReader, Document, VectorStoreIndex, ServiceContext, load_index_from_storage, \
    StorageContext
from llama_index.llms import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.node_parser import SentenceWindowNodeParser
import os
from llama_index.postprocessor import MetadataReplacementPostProcessor, SentenceTransformerRerank

api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"


# 构建句子窗口检索器

# 定义查询引擎函数

# 定义创建向量数据库函数
def build_sentence_window_index(
        documents,
        llm,
        embed_model=OpenAIEmbedding(),
        sentence_window_size=3,
        save_dir="sentence_index",
):
    # create the sentence window node parser w/ default settings
    node_parser = SentenceWindowNodeParser.from_defaults(
        window_size=sentence_window_size,
        window_metadata_key="window",
        original_text_metadata_key="original_text",
    )
    sentence_context = ServiceContext.from_defaults(
        llm=llm,
        embed_model=embed_model,
        node_parser=node_parser,
    )
    if not os.path.exists(save_dir):
        sentence_index = VectorStoreIndex.from_documents(
            documents, service_context=sentence_context
        )
        sentence_index.storage_context.persist(persist_dir=save_dir)
    else:
        sentence_index = load_index_from_storage(
            StorageContext.from_defaults(persist_dir=save_dir),
            service_context=sentence_context,
        )

    return sentence_index


# 定义查询引擎函数
def get_sentence_window_query_engine(
        sentence_index, similarity_top_k=6, rerank_top_n=2
):
    # define postprocessors
    postproc = MetadataReplacementPostProcessor(target_metadata_key="window")
    rerank = SentenceTransformerRerank(
        top_n=rerank_top_n, model="BAAI/bge-reranker-base"
    )

    sentence_window_engine = sentence_index.as_query_engine(
        similarity_top_k=similarity_top_k, node_postprocessors=[postproc, rerank]
    )
    return sentence_window_engine


# 创建向量数据库
index = build_sentence_window_index(
    docs,
    llm=OpenAI(api_key=api_key, model="gpt-3.5-turbo", temperature=0.1),
    save_dir="./sentence_index",
)

# 创业查询引擎
query_engine = get_sentence_window_query_engine(index, similarity_top_k=6)
